[reward, trainer] feat: add managed multi-reward deployments - #524

@chenyingshu PTAL. I've been working on this recently, but it still needs some time to be further refined.

Verdict: the engine path is upstream-aligned, the async API shape is right, and the CPU test suite is genuinely good. Holding merge on three design items (RFC sync for the native backend, manager coupling, native placement addressing) — all are cheap to fix now and expensive to fix after configs pile up on top. The async ask from the earlier thread landed; one concurrency gap below.

Blocking — design

1. Native backend: right call, wrong paper trail — amend RFC #432 and pin the contract now

reward_model_executor.py:65, reward_model.py:227

Native is the right call for this repo: PickScore, HPSv3, CLIP scorers and most small non-LLM reward checkpoints will never be a vLLM/SGLang serving priority. But RFC #432 — the governing design doc this PR implements — says the opposite in three places:

• "The adapter does not implement model execution, batching, scheduling, parameter sharding, or collectives."
• Non-goals: "It does not add a new inference runtime."
• Non-goals: "It does not move model weights into TaskRunner or RewardLoopWorker."

NativeRewardExecutor does all three (weights in reward-loop workers, manager-side batching via even split + padding, framework-side placement). The PR description doesn't mention this divergence either.

The boundary question needs answering deliberately, because the follow-up plan (CPU-native, FSDP) walks further in. Upstream verl's history is the cautionary tale: it deprecated and deleted its framework-hosted FSDP/Megatron reward-model workers (cb01f10ba, #3352; removed in #5255) because sharded model execution inside the trainer was the maintenance cost center — that's the subsystem being incrementally rebuilt if FSDP lands. The property that keeps native healthy is what this PR already does: replica-only, user owns the model class (executor.model), framework owns only placement/lifecycle/dispatch.

Ask: update RFC #432 in this PR to make native a first-class backend, with the contract written down: (a) replica-only — no framework-side TP/PP/FSDP, models too big for one device belong in the engine backend; (b) the user-supplied executor class owns inference; (c) infer(*args, **kwargs) is the deliberately untyped adapter-internal contract (fine — but say so, it's the load-bearing interface every native adapter builds against). If FSDP-native is genuinely wanted later, that's an RFC decision, not a follow-up PR.

2. Named models are locked to the visual manager, via a silent config rewrite

reward_loop.py:129

with open_dict(config.reward.reward_manager):
    config.reward.reward_manager.name = "MultiVisualRewardManager"

Two problems:

The rewrite breaks non-default configs. Upstream resolves managers via resolve_reward_manager_cls on source (register | importlib) + module.path + name. Mutating only .name works solely for verl-omni's exact default: with source: register it raises "Unknown reward manager" (nothing registers that name), and with a user's module.path it tries to load MultiVisualRewardManager from the user's module. The PR body promises "does not silently rewrite legacy configurations" — this is a silent rewrite of the new configuration path.

The lock-in blocks half this repo's use cases. reward.models is presented as general, but scoring flows through MultiVisualRewardManager.run_single, which consumes the visual sample contract (solution_image …). Audio/omni reward models can never use named models as-is. The docs admit it ("Named-model aggregation currently uses the visual reward manager contract") — that's a limitation only if it stays.

The fix is cheap because the abstraction already exists: VisualRewardManager already extends upstream's modality-neutral RewardManagerBase, and upstream's managers (naive, dapo, gdpo, limited, remote) all register under modality-neutral names via @register.

Ask: (a) don't rewrite — validate: error if the configured manager is user-set or not multi-capable; (b) introduce a modality-neutral MultiRewardManager on RewardManagerBase holding the executor dispatch (set_reward_executors) + weighted sum + required policy, with visual sample extraction in a MultiVisualRewardManager(MultiRewardManager) subclass. Then reward_loop.py:129 disappears.

3. placement.devices addressing is order-dependent and unrooted

reward_model.py:140-172, reward.yaml

Native indices address an internal native subpool that is allocated after all engine allocations, in engine config order, and sized max(index)+1. Consequences:

• Adding or reordering one engine model silently remaps what every native index means.
• Sparse indices strand bundles ([0,1] + [4,5] → indices 2,3 allocated, unused; plus the leftover tail pool).
• The docs need repeated disclaimers ("not physical CUDA/NPU IDs, not TP ranks") — when an abstraction needs three paragraphs of disclaimers, that's the design asking to be changed.

Ask: one subpool per native model (exactly what engine models already get), so devices are per-model, stable, and dense; the shared-native-subpool becomes an implementation detail that can vanish. Direction check for the long term: upstream async trainers are moving toward elastic accelerator allocation (dynamic resource scheduling #6556, idle-GPU lending #7373) — keep the static split internal so a future dynamic allocator can replace it without breaking configs.

Should fix

4. async_compute_rm_score is not single-flight

reward_loop.py (async_compute_rm_score), reward_model_executor.py:114

The worker-side drain (lock + inflight + idle event, with the re-check loop in sleep()) is correctly built. But idle-tracking can only see requests that have landed on the worker. With offload: true: batch A finishes and enters its finally → sleep() observes idle while batch B's compute_score_batch RPCs are still in Ray dispatch → A destroys the model → B's infers raise "not awake". Trainers call it serially today, so it's latent — but publishing an async API is an invitation to overlap.

Ask: a manager-level asyncio.Lock around the wake/score/sleep bracket (two lines), which also matches RFC #432's "sequential execution on shared resources" stance until phase-3 admission fencing.

5. Config schemas and trainer predicates live in the wrong layer

reward_loop/reward_model_config.py (392 lines)

Upstream keeps reward schemas in verl/workers/config/reward.py (RewardModelConfig at :82, post-#5255 refactor); verl-omni already mirrors that layout with verl_omni/workers/config/{diffusion,omni}. And the trainer predicates (reward_is_enabled, reward_role_required, reward_pool_is_separate, streaming_reward_enabled) are trainer concerns — upstream's equivalent need_reward_model lives in verl/trainer/ppo/utils.py, not in the reward loop. Right now five trainer modules import from verl_omni.reward_loop.reward_model_config, inverting the layering (trainers reaching into an experimental reward-loop package).

Ask: schemas → verl_omni/workers/config/reward.py; predicates → trainer utils (or the config module's public API), so trainer imports stop crossing into reward_loop.

6. Pooling worker extension is misplaced and untracked

reward_loop/vllm_worker.py:20

Every vLLM colocate extension in verl-omni lives in verl_omni/workers/rollout/vllm_rollout/ — vLLMOmniColocateWorkerExtension (utils.py:35), NPU variant (npu_utils.py:123). A module named vllm_worker.py inside reward_loop/ misleads on both counts (not a reward-loop worker; not where vLLM extensions live). The PR body admits it's temporary ("can be removed after upstream verl handles pooling models generically") but the code carries no TODO(owner) + upstream issue, per repo policy for temporary workarounds. Note upstream already supports pooling RMs elsewhere (/classify determinism, #7027) — the generative assumption in the colocated extension is arguably an upstream bug worth filing rather than shimming forever.

Ask: move next to the other extensions, add TODO(owner) + upstream issue link, and prefer fixing upstream.

7. Dead code in the manager

• reward_loop.py:87 — OmniRewardLoopWorker.close() is never invoked: upstream RewardLoopWorker has no shutdown contract and no trainer path calls it. Resident native models (offload: false) therefore hold accelerator memory until process exit. Wire it into trainer shutdown or drop it.
• reward_loop.py:122 — the second reward.reward_model.enable raise is unreachable (reward_model.py:52 already raised for models+enable, and accelerator-workers+enable raises on the line above). Drop it.

Nits

• reward_loop.py _resolve_result/_gather_results double-wrap Ray refs — ObjectRefs are directly awaitable inside asyncio.gather; one layer is enough.
• MultiRewardModelManager exposes internals (native_device_assignments, native_resource_pool, bind_native_workers) for OmniRewardLoopManager to consume — friend-class coupling across modules. A single create_native_worker_group(name, factory) hook would tighten the surface.
• resolved_offload accepting rollout.free_cache_engine/enable_sleep_mode as aliases gives two spellings for one semantic — your own cleanup plan covers this; make sure it survives the rebase.
• NativeManagedRewardModel._resident is controller-side bookkeeping of worker-side state — fine while these Ray actors don't restart, worth a comment saying so.

PR hygiene

• PR body is stale vs the code: describes reward.deployments while the shipped config is reward.models; links the guide at a fork path (docs/algo/managed_reward_deployments.md) while the real file is docs/algo/named_reward_models.md.
• No test-run evidence in the body — per contribution policy (AI-assisted PRs) add the CPU test command + results and the NPU e2e status.

What landed well (for calibration)

• required/optional failure policy with required=false default preserves legacy zero-contribution semantics — the test migration is honest about the behavior change.
• Engine deployments wrapping verl.RewardModelManager per model is RFC phase 1 exactly as written; ManagedRewardModel's async wake_up/sleep ABC is the thin backend boundary the RFC asked for.
• Async discipline beats upstream: try/finally sleep (upstream skips sleep on failure), sleep() gathers with return_exceptions=True and raises first, reverse-order sleep, sync entry raises inside a running loop instead of blocking it.
• The CPU suite covers schema validation, pool splitting/overcommit, lifecycle delegation, executor contracts, and aggregation failure modes — this is the right test strategy for infra that can't run on CI GPUs.

AI assistance (ZCode) was used for this review.

Thanks for the detailed follow-up review. I addressed the items that belong in this PR and recorded the larger migrations separately in #432.

I also fixed a regression found by the manual NPU run: synchronous scoring could call rollout sleep twice. The compatibility wrapper now delegates lifecycle ownership only to the async path.

Validation for the current head:

• 116 focused CPU tests passed in the verl-omni:gpu container;
• Ruff lint and format checks passed;
• the relocated pooling extension was dynamically imported successfully;
• the engine/native PickScore recipe was manually rerun successfully on the 16-NPU setup.

Deferred follow-up work is tracked in the table on #432: unified MultiRewardManager, Audio/Visual manager consolidation, legacy path deprecation after model migrations, CPU-native execution, native FSDP, vLLM-Omni serving, dynamic load balancing, and full async/streaming integration.

@zhtmike CI passed